//      svm.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include "svm.h"

KernelCache::KernelCache(int data_size, boost::program_options::variables_map& vm)
  : time_stamp_(data_size, 0), now_(1)
{
  size_ = vm.count("kernel_cache") ? vm["kernel_cache"].as<int>() 
    : (1024*1024*1024 / (sizeof(double) * data_size));
  MSG("kernel cache size="<<size_);
}

void KernelCache::add(int i, const arma::vec& k)
{
  c_[i] = k;
  time_stamp_[i] = now_++;

  typedef std::map<int, arma::vec>::value_type value_type;
  if (sz(c_) > size_) { // remove LRU item    
    int min_i = -1;
    long long min = now_+1;
    BOOST_FOREACH(const value_type& cache_item, c_) {
      int j = cache_item.first;
      if (time_stamp_[j] < min) {
	min = time_stamp_[j];
	min_i = j;
      }
    }
    c_.erase(min_i);
  }
}

bool KernelCache::contains(int i, int j)
{
  return i == j || c_.find(i) != c_.end() || c_.find(j) != c_.end();
}

double KernelCache::kval(int i, int j)
{
  if (i == j) return diag_(i);
  typedef std::map<int, arma::vec>::iterator iterator;
  iterator i_iter = c_.find(i);
  if (i_iter != c_.end()) {
    time_stamp_[i] = now_++;
    return i_iter->second(j);
  }
  else {
    iterator j_iter = c_.find(j);
    if (j_iter != c_.end()) {
      time_stamp_[j] = now_++;
      return j_iter->second(i);
    }
    else
      FATAL_ERROR("Cannot find cache element i=" << i << " j=" << j);
  }
}

// ====================== Kernel =========================

Kernel::Kernel(MatrixData& data, boost::program_options::variables_map& vm)
  : data_(data), vm_(vm), type_(LINEAR), cache_(sz(data), vm)
{
  if (vm_.count("kernel")) {
    std::string kernel = vm_["kernel"].as<std::string>();
    if (kernel == "0" || kernel == "linear") 
      type_ = LINEAR;
    if (kernel == "1" || kernel == "gaussian") 
      type_ = GAUSSIAN;
    else
      FATAL_ERROR("Unknown kernel type = " << kernel);
  }

  MSG("kernel type="<<type_);

  switch (type_) {
  case LINEAR:
    kernel_function_ = linear_kernel;
    break;
  case GAUSSIAN:
    kernel_function_ = gaussian_kernel;
    break;
  default:
    FATAL_ERROR("Unknown kernel type = " << type_);
  }

  arma::vec diag(sz(data));
  for (int i = 0; i < sz(data); i++) diag(i) = kernel_function_(data_.x(i), data_.x(i));
  cache_.set_diag(diag);
}

double Kernel::operator() (int i, int j)
{
  if (cache_.contains(i, j))
    return cache_.kval(i, j);
  else {
    arma::vec k = this->operator()(j);
    cache_.add(j, k);
    return k(i);
  }
}

double Kernel::operator() (int i, const arma::vec& x)
{
  return kernel_function_(data_.x(i), x);
}

arma::vec Kernel::operator() (int j)
{
  arma::vec k = arma::vec(sz(data_));
  for (int i = 0; i < sz(data_); i++)
    k(i) = kernel_function_(data_.x(i), data_.x(j));
  return k;
}

arma::vec Kernel::operator() (const arma::vec& x)
{
  arma::vec k = arma::vec(sz(data_));
  for (int i = 0; i < sz(data_); i++)
    k(i) = this->operator()(i, x);
  return k;
}

double Kernel::linear_kernel(const arma::vec& x, const arma::vec& y)
{
  return arma::dot(x, y);
}

double Kernel::gaussian_kernel(const arma::vec& x, const arma::vec& y)
{
  return exp(-0.5*arma::norm(x-y,2));
}

// ================  L1SVM  ====================
KernelL1SVM::KernelL1SVM(MatrixData& data, boost::program_options::variables_map& vm, Kernel& K) 
  : data_(data), vm_(vm), K_(K)
{
  //MSG("data=" << data_.size() << " x " << data_.dim());
}


void KernelL1SVM::solve(bool use_bias) 
{
  p_.alpha = arma::zeros<arma::vec>(sz(data_));
  p_.bias = 0;
  p_.sp.clear();
  g_ = arma::ones<arma::vec>(sz(data_));

  int max_iter = vm_.count("optim_iter") ? vm_["optim_iter"].as<int>() : 100*sqr(sz(data_));
  double tol = vm_.count("optim_tol") ? vm_["optim_tol"].as<double>() : 1e-8;
  C_ = vm_.count("C") ? vm_["C"].as<double>() : 100;

  bool termination_reached = false;
  for (int iter = 0; iter < max_iter; iter++) {
    double Gmax = -INFINITY, Gmin = INFINITY;
    int i = -1, j = -1;
    
    for (int i1 = 0; i1 < sz(data_); i1++) {
      double y_i = data_.y(i1);
      double g_i = g_(i1);
      double a_i = p_.alpha(i1);
      if ((y_i < 0 && a_i > 0) || (y_i > 0 && a_i < C_))
	if (y_i * g_i > Gmax) {
	  Gmax = y_i * g_i;
	  i = i1;
	}
    }

    /*
    for (int j1 = 0; j1 < sz(data_); j1++) {
      double y_j = data_.y(j1);
      double g_j = g_(j1);
      double a_j = p_.alpha(j1);
      if ((y_j < 0 && a_j < C_) || (y_j > 0 && a_j > 0))
	if (y_j*g_j < Gmin) {
	  Gmin = y_j*g_j;
	  j = j1;
	}
    }
    */

    double max = -INFINITY;
    for (int j1 = 0; j1 < sz(data_); j1++) {
      double y_j = data_.y(j1);
      double g_j = g_(j1);
      double a_j = p_.alpha(j1);
      double G = y_j*g_j;
      if (G < Gmax && ((y_j < 0 && a_j < C_) || (y_j > 0 && a_j > 0))) {
	double s = 0.5*sqr(Gmax-G)/(K_(i,i)+K_(j1,j1)-2*K_(j1,i));
	if (s > max) {
	  Gmin = G;
	  j = j1;
	  max = s;
	}
      }
    }
    

    if (iter % 1000 == 0)
    MSG_VERBOSE(vm_, "iter=" << iter << " Gmax-Gmin=" << Gmax-Gmin << " i=" << i << " j=" << j);

    if (Gmax <= Gmin+tol) { // termination condition
      p_.bias = Gmax;
      termination_reached = true;
      MSG_VERBOSE(vm_, "termination condition reached Gmax=" << Gmax << " <= Gmin=" << Gmin);
      break;
    }

    double a_i = p_.alpha(i), a_j = p_.alpha(j);
    double y_i = data_.y(i), y_j = data_.y(j);
    double lambda = (y_i > 0 ? C_ : 0) - y_i*a_i;
    double lambda1 = y_j*a_j - (y_j > 0 ? 0 : -C_);
    double lambda2 = (Gmax-Gmin) / (K_(i,i)+K_(j,j)-2*K_(j,i));
    if (lambda > lambda1) lambda = lambda1;
    if (lambda > lambda2) lambda = lambda2;
    if (iter % 1000 == 0) MSG_VERBOSE(vm_, "lambda="<<lambda);
    
    //update gradient
    for (int k = 0; k < sz(data_); k++) {
      double y_k = data_.y(k);
      g_(k) = g_(k) - lambda*y_k*K_(k,i) + lambda*y_k*K_(k,j);
    }
    //update alpha
    p_.alpha(i) += y_i*lambda;
    p_.alpha(j) -= y_j*lambda;
    
    if (p_.alpha(i) < 0 || p_.alpha(i) > C_ || p_.alpha(j) < 0 || p_.alpha(j) > C_) {
      FATAL_ERROR("off bound");
    }

    if (p_.alpha(i) < 0) p_.alpha(i) = 0;
    if (p_.alpha(i) > C_) p_.alpha(i) = C_;
    if (p_.alpha(j) < 0) p_.alpha(j) = 0;
    if (p_.alpha(j) > C_) p_.alpha(j) = C_;
  }

  if (!termination_reached) {
    MSG("termination not reached after max_iter=" << max_iter << " try set --optim_iter");
    double n_free = 0, s = 0; 
    for (int i = 0; i < sz(data_); i++) {
      double a_i = p_.alpha(i), y_i = data_.y(i);
      if (a_i > 0 && a_i < C_) {
	n_free += 1.0;
	s += y_i* ( g_(i) - 1 );
      }
    }
    if (n_free > 0)
      p_.bias = s / n_free;
    else {
      p_.bias = 0; // TODO
      MSG("No free support vector");
    }
  }

  if (vm_.count("debug")) {
    print(vm_["debug"].as<std::string>());		
  }
}

KernelL1SVM::Param KernelL1SVM::get_param() const 
{
  return p_;
}

void KernelL1SVM::set_param(const Param& p)
{
  p_ = p;
}

KernelL1SVM::Output KernelL1SVM::output(Data& test) const
{
  arma::vec r = arma::vec(sz(test));
  for (int i = 0; i < sz(test); i++) {
    arma::vec k = K_(test.x(i)); // K( . , x)
    k %= data_.label() % p_.alpha;
    r(i) = arma::sum(k) + p_.bias;
  }
  return r;
}

void KernelL1SVM::print(const std::string& outputFile) const
{
  FileAccessWrite(out, outputFile);
  for (int i = 0; i < sz(p_.alpha); i++)
    out << p_.alpha(i) << " ";
  out << ENDL;
  out << p_.bias << ENDL;
  for (int i = 0; i < sz(p_.sp); i++)
    out << p_.sp[i] << " ";
  out << ENDL;
  out.close();
}

void KernelL1SVM::print_stats(std::ostream& output, const ClassPerf& cp)
{
  output << cp.accuracy() << " " << cp.sensitivity() << " " 
	 << cp.specificity() << " " << cp.mse() << " ";
}

